#!/usr/bin/python

from openeye.oechem import *


import sys, string, os

ofs = oemolostream()
mol = OEGraphMol()

#-------------------------------
def mol2out(mol,dir):

  
  ofs.open(dir + mol.GetTitle() + '.mol2')
  flags = OEOFlavor_PDB_CHARGE 
  OEWriteMolecule(ofs, mol)
#  OEWritePDBFile(ofs, mol, flags)
#  OEWriteXYZFile(ofs,mol)
#-------------------------------
def cal_charge(mol):
  charge = 0

  for atom in mol.GetAtoms():
    charge = charge + atom.GetPartialCharge()
  
  charge = round(charge,0)
  if charge == -0.00:
    charge = 0.00
  #print charge
  return charge
#-------------------------------
def solvent_header(out_file, solvent):

  if solvent =='PENT':
    out_file.write('SOLVNT=GENORG IOFR=1.3575 ALPHA=0.00 BETA=0.00 GAMMA=22.3' + '\n')
    out_file.write('& DIELEC=1.84 FACARB=0.00 FEHALO=0.00 ')
  elif solvent =='PFB':
    out_file.write('SOLVNT=GENORG IOFR=1.3777 ALPHA=0.00 BETA=0.00 GAMMA=31.7' + '\n')
    out_file.write('& DIELEC=2.03 FACARB=0.5 FEHALO=0.50 ')
  elif solvent =='CYC':
    #BWQ definitions, but it is actually hexadecane
    out_file.write('SOLVNT=GENORG IOFR=1.4345 ALPHA=0.00 BETA=0.00 GAMMA=38.93' + '\n')
    out_file.write('& DIELEC=2.06 FACARB=0.00 FEHALO=0.00 ')
  elif solvent =='RCYC':
    #BWQ really cyclohexadecane
    out_file.write('SOLVNT=GENORG IOFR=1.4266 ALPHA=0.00 BETA=0.00 GAMMA=35.5' + '\n')
    out_file.write('& DIELEC=2.02 FACARB=0.00 FEHALO=0.00 ')
  elif solvent =='TMB':
    out_file.write('SOLVNT=GENORG IOFR=1.5048 ALPHA=0.00 BETA=0.19 GAMMA=42.0' + '\n')
    out_file.write('& DIELEC=2.37 FACARB=0.67 FEHALO=0.00 ')
  elif solvent == 'DBE':
    out_file.write('SOLVNT=GENORG IOFR=1.3992 ALPHA=0.00 BETA=0.45 GAMMA=32.3' + '\n')
    out_file.write('& DIELEC=3.05 FACARB=0.00 FEHALO=0.00 ' ) 
  elif solvent =='DEE':
    out_file.write('SOLVNT=GENORG IOFR=1.3526 ALPHA=0.00 BETA=0.41 GAMMA=24.0' + '\n')
    out_file.write('& DIELEC=4.24 FACARB=0.00 FEHALO=0.00 ' ) 
  elif solvent =='EAC':
    out_file.write('SOLVNT=GENORG IOFR=1.3723 ALPHA=0.00 BETA=0.45 GAMMA=33.7' + '\n')
    out_file.write('& DIELEC=5.99 FACARB=0.00 FEHALO=0.00 ' ) 
  elif solvent =='DCA':
    out_file.write('SOLVNT=GENORG IOFR=1.4372 ALPHA=0.37 BETA=0.48 GAMMA=41.0' + '\n')
    out_file.write('& DIELEC=7.53 FACARB=0.00 FEHALO=0.00 ' ) 
  elif solvent =='DCE':
    out_file.write('SOLVNT=GENORG IOFR=1.4448 ALPHA=0.10 BETA=0.11 GAMMA=45.9' + '\n')
    out_file.write('& DIELEC=10.19 FACARB=0.00 FEHALO=0.50 ' ) 
  elif solvent =='TBU':
    out_file.write('SOLVNT=GENORG IOFR=1.3878 ALPHA=0.31 BETA=0.60 GAMMA=28.7' + '\n')
    out_file.write('& DIELEC=12.47 FACARB=0.00 FEHALO=0.00 ' ) 
  elif solvent =='ACE':
    out_file.write('SOLVNT=GENORG IOFR=1.3588 ALPHA=0.04 BETA=0.49 GAMMA=23.5' + '\n')
    out_file.write('& DIELEC=20.49 FACARB=0.00 FEHALO=0.00 ' ) 
  elif solvent =='MET':
    out_file.write('SOLVNT=GENORG IOFR=1.3288 ALPHA=0.43 BETA=0.47 GAMMA=22.1' + '\n')
    out_file.write('& DIELEC=32.63 FACARB=0.00 FEHALO=0.00 ' ) 
  elif solvent =='DMSO':
    out_file.write('SOLVNT=GENORG IOFR=1.4170 ALPHA=0.00 BETA=0.88 GAMMA=61.8' + '\n')
    out_file.write('& DIELEC=46.83 FACARB=0.00 FEHALO=0.00 ' )
  else:
    print "Solvent Keyword is wrong" 

 
#-------------------------------
def header(solvent, filename, out_file, charge):

  if solvent == 'Water':
     out_file.write('\n')
     out_file.write('SOLVNT=WATER CHARGE=' +  str(charge) + ' AM1 SM5.42R 1SCF GEO-OK CART' + '\n')
  else:
    out_file.write('\n')
    solvent_header(out_file, solvent)
    out_file.write('CHARGE=' + str(int(charge)) + ' AM1 SM5.42R 1SCF' + '\n')
    out_file.write('& GEO-OK CART' + '\n' )

#-------------------------------
def convert2amsol(filename, charge, mol, solvent, dir):
  
  global dict

  out_file_wat = open('temp.in-wat', 'w')
  out_file_hex = open('temp.in-hex', 'w')

  command = 'babel -imol2 ' +  dir + filename + '.mol2 -omopcrt '+ dir + filename + '.com'
  #print command
  os.system(command)
  in_file = open(dir + filename + '.com', 'r')
#  out_file = open(filename + '_mod.com', 'w')
  #cut of firt 7 lines
  counter = 0
  while counter < 3:
    in_file.readline()
    counter = counter + 1
  
  header('Water', filename, out_file_wat, charge)  
  if solvent <> 'Vac':
    header(solvent, filename, out_file_hex, charge)
  block = in_file.read()
  
  out_file_wat.write(mol.GetTitle() + ' ' + str(mol.NumAtoms()) + '\n\n')
  if solvent <> 'Vac':
    out_file_hex.write(mol.GetTitle() + ' ' + str(mol.NumAtoms()) + '\n\n')
  

  out_file_wat.write(block)
  if solvent <> 'Vac':
    out_file_hex.write(block)
  
    out_file_hex.write('\n')
  
  out_file_wat.write('\n')
  
 
  
#  out_file.close()
  
  in_file.close()  
  
    
  out_file_wat.close()
  out_file_hex.close()
 
   

#-------------------------------
def Update_charge(mol, partial_charge_list, ofs_new_mol2, update_error):
  
  counter = 0
    
  if len(partial_charge_list) == mol.NumAtoms():
  
    for atom in mol.GetAtoms():
      atom.SetPartialCharge(partial_charge_list[counter])
      counter = counter + 1 
    
    OEWriteMolecule(ofs_new_mol2, mol)
    #the format is bad, but that"s all you can get with these unfinished OE tools
  else:
    update_error.write(mol.GetTitle() + '\n')
  
#-------------------------------
def run_amsol():

   os.system('amsol-mod4 temp.in-wat temp.o-wat')
   os.system('amsol-mod4 temp.in-hex temp.o-hex')
     
   
#-------------------------------
def clean_up(dir, moli):
   print 'cleaning up'
   command = 'rm -rf ' + dir
   
   os.system(command)

   #command = 'rm -rf fort.* temp.o*'
   print 'clean up'
   os.system(command)
   os.system('mv tab.solv ' + moli +'.solv')
   os.system('mv ' + moli + ' ' + moli + '.nmol2')
#-------------------------------

def parse_results(filename, mol, charge, result_tab, ofs_new_mol2, update_error,cav_term, debug):

  tab_list = []

  result_file_wat = open('temp.o-wat', 'r')
       
  part_pol_sol_wat = []
  part_apol_sol_wat = []
  
  energy = 'false'
  start = 'false'
  almost_there = 'false'
 
  try:
       
       
       
       for i in result_file_wat.xreadlines():
         
	 if energy == 'true':
	   pol_sol_wat = float(i[30:38])
	   apol_sol_wat =float(i[51:59])
	   energy = 'false'
	   if debug == 'true':
                print 'pol_sol_wat', pol_sol_wat
                print 'apol_sol_wat', apol_sol_wat
       
         if start == 'true':
	   try:
	     part_pol_sol_wat.append(float(i[30:38]))
	     part_apol_sol_wat.append(float(i[51:59]))
	   except:
	     start = 'false'
	     energy = 'true'
	 
	 if almost_there == 'true':
	   start = 'true'
	   almost_there = 'wrong'
         
	 if i[62:68] == '(kcal)':
	   #data starts in line after following line
	   almost_there ='true'
	   #print ' =====> test'
	   
       result_file_wat.close()
	   
	   
       result_file_hex = open('temp.o-hex', 'r')
       
       part_pol_sol_hex = []
       part_apol_sol_hex = []
       charge_list = []
       
        
       for i in result_file_hex.xreadlines():
         
	 if i[6:13] == 'G-P(sol':
	   pol_sol_hex = float(i[65:74])
           if debug == 'true':
                print 'pol_sol_hex', pol_sol_hex
	 elif i[6:16] =='G-CDS(sol)':
	   apol_sol_hex =float(i[65:74])
           if debug == 'true':
                print 'apol_sol_hex', apol_sol_hex     
      
         if start == 'true':
	   try:
	     part_pol_sol_hex.append(float(i[30:38]))
	     part_apol_sol_hex.append(float(i[51:59]))
	     if debug == 'true':
                 print i[51:59]    
	     charge_list.append(float(i[18:26]))
	   except:
	     start = 'false'
	 
	 if almost_there == 'true':
	   start = 'true'
	   almost_there = 'wrong'
         
	 if i[62:68] == '(kcal)':
	   #data starts in line after following line
	   almost_there ='true'
	   #print ' =====> test'
	   
	 #get correction_factor for non-water solvents
	 if i[4:16] == 'Contribution':
		cs_contribution = float(i[52:59])
		#print '------------->  test'

       result_file_hex.close()

 
       
       #everything is read in
       desolv_pol = pol_sol_wat - pol_sol_hex
       if cav_term == 'cav':
         desolv_apol = apol_sol_wat - apol_sol_hex
       else:
         desolv_apol = apol_sol_wat

       correction_factor = cs_contribution/len(charge_list)
       if debug == 'true':
		print 'Corrrection factor :' + str(correction_factor)
       

       
       
       if cav_term <> 'sep':	
         line = filename + string.rjust(str(mol.NumAtoms()), 4)  + '%5.1f'%charge + '%9.2f'%desolv_pol + '   000.00' +'%9.2f'%desolv_apol + '   -00.00' + '\n'
       else:  
	 line = filename + '\t' + string.rjust(str(mol.NumAtoms()), 4)   + '\t' + '%8.2f'%charge + '\t' + '%8.2f'%pol_sol_wat + '\t' + '%8.2f'%pol_sol_hex +'\t' + '%8.2f'%apol_sol_wat + '\t' + '%8.2f'%apol_sol_hex + '\n'
 	 #line = filename + '\t' + int(string.rjust(str(mol.NumAtoms()), 4))  + '\t' + '%8.2f'%charge + '\t' + '%8.2f'%pol_sol_wat + '\t' + '%8.2f'%pol_sol_hex +'\t' + '%8.2f'%apol_sol_wat + '\t' + '%8.2f'%apol_sol_hex + '\n'
       tab_list.append(line)


       counter = 0
       for charge in charge_list:
         #print charge  
         if cav_term == 'cav':
	   #correct apol_sol for cs_contribution
      	   tab_list.append('%8.4f'%charge + '%8.2f'%(part_pol_sol_wat[counter] - part_pol_sol_hex[counter]) + '  00.30' + '%8.2f'%((part_apol_sol_wat[counter]) - (part_apol_sol_hex[counter] + correction_factor)) + '   -0.30' + '\n')
         elif cav_term == 'sep':
      	   tab_list.append('%8.4f'%charge + '\t' + '%8.2f'%part_pol_sol_wat[counter] + '\t' + '%8.2f'%part_pol_sol_hex[counter] + '\t' + '%8.2f'%part_apol_sol_wat[counter] + '\t' + '%8.2f'%part_apol_sol_hex[counter] + '\n')
      	 else:
	   tab_list.append('%8.4f'%charge + '%8.2f'%(part_pol_sol_wat[counter] - part_pol_sol_hex[counter]) + '  00.30' + '%8.2f'%(part_apol_sol_wat[counter]) + '   -0.30' + '\n')
	 counter = counter + 1
       Update_charge(mol,charge_list,ofs_new_mol2, update_error)
       
       for lines in tab_list:
      	result_tab.write(lines)
  
       
       
  except:
    update_error.write(mol.GetTitle() + '\n')
  if debug <> 'true':
      os.system('rm -f temp.o*')

#---------------------------
def parse_results_vac(filename, mol, charge, result_tab, ofs_new_mol2, update_error,cav_term):

  tab_list = []

  result_file_wat = open('temp.o-wat', 'r')
       
  part_pol_sol_wat = []
  part_apol_sol_wat = []
  charge_list = []
  
  energy = 'false'
  start = 'false'
  almost_there = 'false'
 
  try:
       
       
       for i in result_file_wat.xreadlines():
         
	 if energy == 'true':
	   pol_sol_wat = float(i[30:38])
	   apol_sol_wat =float(i[51:59])
	   energy = 'false'
       
         if start == 'true':
	   try:
	     part_pol_sol_wat.append(float(i[30:38]))
	     part_apol_sol_wat.append(float(i[51:59]))
	     charge_list.append(float(i[18:26]))
	   except:
	     start = 'false'
	     energy = 'true'
	 
	 if almost_there == 'true':
	   start = 'true'
	   almost_there = 'wrong'
         
	 if i[62:68] == '(kcal)':
	   #data starts in line after following line
	   almost_there ='true'
	   #print ' =====> test'
	   
       result_file_wat.close()
	   
	   
       result_file_hex = open('temp.o-hex', 'r')
       
       part_pol_sol_hex = []
       part_apol_sol_hex = []
       charge_list = []
       
       
     
       #everything is read in
       desolv_pol = pol_sol_wat 
       desolv_apol = apol_sol_wat
       
       if cav_term <> 'sep':	
         line = filename + string.rjust(str(mol.NumAtoms()), 4)  + '%5.1f'%charge + '%9.2f'%desolv_pol + '   000.00' +'%9.2f'%desolv_apol + '   -00.00' + '\n'
       else:  
	 line = filename + '\t' + string.rjust(str(mol.NumAtoms()), 4)   + '\t' + '%8.2f'%charge + '\t' + '%8.2f'%pol_sol_wat + '\t' + '%8.2f'%pol_sol_hex +'\t' + '%8.2f'%apol_sol_wat + '\t' + '%8.2f'%apol_sol_hex + '\n'
 	 #line = filename + '\t' + int(string.rjust(str(mol.NumAtoms()), 4))  + '\t' + '%8.2f'%charge + '\t' + '%8.2f'%pol_sol_wat + '\t' + '%8.2f'%pol_sol_hex +'\t' + '%8.2f'%apol_sol_wat + '\t' + '%8.2f'%apol_sol_hex + '\n'
       tab_list.append(line)

       counter = 0

       for charge in charge_list:
         if cav_term == 'cav':
     	   tab_list.append('%8.4f'%charge + '%8.2f'%(part_pol_sol_wat[counter]) + '  00.30' + '%8.2f'%(part_apol_sol_wat[counter]) + '   -0.30' + '\n')
         elif cav_term == 'sep':
      	   tab_list.append('%8.4f'%charge + '\t' + '%8.2f'%part_pol_sol_wat[counter] + '\t' + '%8.2f'%part_pol_sol_hex[counter] + '\t' + '%8.2f'%part_apol_sol_wat[counter] + '\t' + '%8.2f'%part_apol_sol_hex[counter] + '\n')
      	 else:
	   tab_list.append('%8.4f'%charge + '%8.2f'%(part_pol_sol_wat[counter]) + '  00.30' + '%8.2f'%(part_apol_sol_wat[counter]) + '   -0.30' + '\n')
	 counter = counter + 1
       Update_charge(mol,charge_list,ofs_new_mol2, update_error)
       
       for lines in tab_list:
      	result_tab.write(lines)
  
       
       
  except:
    update_error.write(mol.GetTitle() + '\n')

  os.system('rm -f temp.o*')
